import array
import struct

from parse.resolver import VariableType
from vm.chunk import Chunk
from vm.value import Value, ValueType, Obj, ObjInstance, ObjClosure, ObjType
from vm.instruction import OpCode
from typing import Dict, List, Optional

STACK_SIZE = 1024


class VM:
    def __init__(self):
        self.strings: Dict[str, 'Obj'] = {}
        self.globals: Dict[str, Value] = {}
        self.objects: List[Obj] = []
        self.debug = True
        self.frames: List['CallFrame'] = []

    def push(self, value: Value):
        self.stack.append(value)

    def pop(self) -> Value:
        return self.stack.pop()

    def pop_pair(self) -> (Value, Value):
        second = self.pop()
        first = self.pop()

        return first, second

    def peek(self, offset=0):
        index = len(self.stack) - 1 - offset
        return self.stack[index]

    def copy_string(self, s: str) -> Obj:
        if s in self.strings.keys():
            return self.strings[s]
        else:
            obj = Obj.string(s)
            self.strings[s] = obj
            return obj

    def read_byte(self):
        ret = self.chunk.code[self.ip]
        self.ip += 1
        return ret

    def read_constant(self) -> Value:
        index = self.read_byte()

        return self.chunk.values[index]

    def current_call_frame(self):
        return self.frames[-1]

    @property
    def ip(self):
        return self.current_call_frame().ip

    @ip.setter
    def ip(self, v: int):
        self.current_call_frame().ip = v

    @property
    def chunk(self):
        return self.frames[-1].closure.fun.chunk

    @property
    def stack(self):
        return self.frames[-1].stack

    def run(self):
        while True:
            if self.debug:
                self.show_trace_info()

            opcode = self.read_byte()
            # print(f'opcode = {OpCode(opcode).name}')
            match opcode:
                case OpCode.RETURN:
                    should_exit = self.finish_call_frame()
                    if should_exit:
                        break
                case OpCode.CONSTANT:
                    self.push(self.read_constant())
                case OpCode.NEGATE:
                    num_value = self.pop().number_value()
                    self.push(Value.new_number(num_value * (-1)))
                case OpCode.ADD:
                    self.add()
                case OpCode.SUBTRACT:
                    self.sub()
                case OpCode.MULTIPLY:
                    self.mul()
                case OpCode.DIVIDE:
                    self.div()
                case OpCode.NOT:
                    v = self.pop()
                    self.push(Value.boolean(not bool(v)))
                case OpCode.NIL:
                    self.push(Value.nil())
                case OpCode.TRUE:
                    self.push(Value.boolean(True))
                case OpCode.FALSE:
                    self.push(Value.boolean(False))
                case OpCode.EQUAL:
                    arg1, arg2 = self.pop_pair()
                    self.push(Value.boolean(arg1 == arg2))
                case OpCode.GREATER:
                    arg1, arg2 = self.pop_pair()
                    self.push(Value.boolean(arg1.number_value() > arg2.number_value()))
                case OpCode.LESS:
                    arg1, arg2 = self.pop_pair()
                    self.push(Value.boolean(arg1.number_value() < arg2.number_value()))
                case OpCode.PRINT:
                    print(self.peek())
                case OpCode.POP:
                    self.pop()
                case OpCode.GET_LOCAL:
                    offset = self.read_byte()
                    self.push(self.stack[offset])
                case OpCode.SET_LOCAL:
                    offset = self.read_byte()
                    self.stack[offset] = self.peek()
                case OpCode.DEFINE_GLOBAL:
                    self.globals[self.read_constant().string_value()] = self.peek()
                case OpCode.GET_GLOBAL:
                    self.push(self.get_global())
                case OpCode.SET_GLOBAL:
                    self.globals[self.read_constant().string_value()] = self.peek()
                case OpCode.GET_UPVALUE:
                    self.push(self.get_upvalue())
                case OpCode.SET_UPVALUE:
                    self.set_upvalue()
                case OpCode.GET_FREE_VALUE:
                    self.push(self.get_free_value())
                case OpCode.SET_FREE_VALUE:
                    self.set_free_value()
                case OpCode.GET_PROPERTY:
                    self.handle_get_property()
                case OpCode.SET_PROPERTY:
                    self.hand_set_property()
                case OpCode.GET_SUPER:
                    self.hand_super()
                case OpCode.ESCAPE_VALUE:
                    v = self.pop()
                    value = Value.wrap(v)
                    self.add_obj(value.value)
                    self.push(value)
                case OpCode.JUMP:
                    offset_bytes = [self.read_byte(), self.read_byte()]
                    offset, *_useless = struct.unpack('>h', array.array('B', offset_bytes))
                    self.ip += offset
                case OpCode.JUMP_IF_FALSE:
                    if not bool(self.peek()):
                        offset_bytes = [self.read_byte(), self.read_byte()]
                        offset, *_useless = struct.unpack('>h', array.array('B', offset_bytes))
                        self.ip += offset
                    else:
                        self.ip += 2
                case OpCode.SUPER_INVOKE:
                    raise NotImplementedError()
                case OpCode.CLOSURE:
                    self.closure()
                case OpCode.CLASS:
                    # 创建一个类到栈顶
                    self.handle_class()
                case OpCode.INHERIT:
                    self.inherit()
                case OpCode.METHOD:
                    # 将栈顶的闭包对象放入其下方的类对象中
                    self.handle_method()
                case OpCode.CALL:
                    self.handle_call()
                case _:
                    raise NotImplementedError(f'Op code {opcode} not recognized')

    def show_trace_info(self):
        print("          ", end='')
        for v in self.stack_values():
            print("[ {} ]".format(v))
        print()

        self.chunk.disassemble_instruction(self.ip)

    def add(self):
        arg1, arg2 = self.pop_pair()
        result = arg1.number_value() + arg2.number_value()
        self.push(Value.new_number(result))

    def sub(self):
        arg1, arg2 = self.pop_pair()
        result = arg1.number_value() - arg2.number_value()
        self.push(Value.new_number(result))

    def mul(self):
        arg1, arg2 = self.pop_pair()
        result = arg1.number_value() * arg2.number_value()
        self.push(Value.new_number(result))

    def stack_values(self):
        return iter(self.stack)

    def div(self):
        arg1, arg2 = self.pop_pair()
        result = arg1.number_value() / arg2.number_value()
        self.push(Value.new_number(result))

    def add_obj(self, v: Obj):
        self.objects.append(v)

    def handle_call(self):
        arg_count = self.read_byte()

        args = []

        for i in range(arg_count):
            args.append(self.pop())

        args.reverse()

        callee = self.pop()

        if callee.value_type != ValueType.OBJ:
            raise ValueError(f'invalid callee type {callee.value_type}')

        if callee.value.obj_type == ObjType.CLOSURE:
            closure = callee.closure_value()

            if arg_count != closure.fun.arity:
                raise ValueError(f'invalid argument count {arg_count}, should be {closure.fun.arity}')

            call_frame = CallFrame(closure, args)
            self.frames.append(call_frame)
        elif callee.value.obj_type == ObjType.CLASS:
            klass = callee.class_value()
            # todo: 添加初始化方法的调用
            self.push(Value(ValueType.OBJ, Obj.instance(klass)))
        elif callee.value.obj_type == ObjType.BOUND_METHOD:
            method = callee.value.bound_method_value()
            arity = method.method.fun.arity
            args = [method.receiver] + args
            if arg_count != arity:
                raise ValueError(f'invalid argument count {arg_count}, should be {arity}')
            call_frame = CallFrame(method.method, args)
            self.frames.append(call_frame)

    def finish_call_frame(self):
        """
        :return: 是否需要退出程序
        """
        if len(self.frames) == 1:
            return True

        ret_val = self.pop()

        self.frames.pop()
        self.push(ret_val)

    def get_global(self):
        name = self.read_constant().string_value()
        if name not in self.globals:
            raise ValueError(f'cannot get global value "{name}"')
        return self.globals[name]

    def get_upvalue(self):
        offset = self.read_byte()
        return self.current_call_frame().closure.upvalues[offset].unwrap()

    def set_upvalue(self):
        offset = self.read_byte()
        self.current_call_frame().closure.upvalues[offset].update_upvalue(self.peek())

    def get_free_value(self):
        offset = self.read_byte()
        return self.stack[offset].unwrap()

    def set_free_value(self):
        offset = self.read_byte()
        self.stack[offset].update_upvalue(self.peek())

    def closure(self):
        value = self.read_constant()
        is_fun = value.value_type == ValueType.OBJ and value.value.obj_type == ObjType.FUNCTION
        if not is_fun:
            raise TypeError(f'should be a function, but {value}')

        fun = value.value.func_value()

        upvalues = []
        for _ in range(fun.upvalue_count):
            value_type = VariableType(self.read_byte())
            offset = self.read_byte()
            if value_type == VariableType.UPVALUE:
                upvalues.append(self.current_call_frame().closure.upvalues[offset])
            elif value_type == VariableType.FREE:
                upvalues.append(self.current_call_frame().stack[offset])

        closure = Value.closure(fun, upvalues)
        self.push(closure)

    def handle_class(self):
        cls_name = self.read_constant().string_value()
        self.push(Value.klass(cls_name))

    def handle_method(self):
        method = self.peek().closure_value()
        self.peek(1).value.class_value().add_method(method.fun.name, method)

    def handle_get_property(self):
        instance = self.peek().instance_value()
        name = self.read_constant().string_value()
        if name in instance.fields:
            self.pop()
            self.push(instance.fields[name])
        elif name in instance.klass.methods:
            method = instance.klass.methods[name]
            bound_method = Obj.bound_method(self.pop(), method)
            self.push(Value(ValueType.OBJ, bound_method))

    def hand_set_property(self):
        property_name = self.read_constant().string_value()
        obj = self.peek(1).instance_value()
        obj.fields[property_name] = self.peek()
        value = self.pop()
        self.pop()
        self.push(value)

    def inherit(self):
        super_cls = self.peek().class_value()
        cls = self.peek(1).class_value()
        cls.methods.update(super_cls.methods)

    def hand_super(self):
        super_cls = self.pop().class_value()
        method_name = self.read_constant().string_value()
        method = super_cls.methods[method_name]
        bound_method = Obj.bound_method(self.pop(), method)
        self.push(Value(ValueType.OBJ, bound_method))


class CallFrame:
    def __init__(self, closure: ObjClosure, stack: List[Value]):
        self.ip = 0
        self.closure = closure
        self.stack = stack
